import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

def plot_results(df):
    """
    Reads the CSV data and generates enhanced performance plots.
    """
    sns.set_theme(style="darkgrid")

    # 筛选每个实验的数据
    df_filesize = df[df['Experiment'] == 'FileSize'].copy()
    df_cspcount = df[df['Experiment'] == 'CSPCount'].copy()
    df_numblocks = df[df['Experiment'] == 'NumBlocks'].copy() 
    # 按块数量排序，因为块大小增加时，块数量是减少的
    df_numblocks = df_numblocks.sort_values(by='Value').reset_index(drop=True)

    y_columns = ['EncryptTimeMs', 'AuditTimeMs', 'ProofGenTimeMs', 'VerifyTimeMs']
    legend_labels = ['Encrypt (DO)', 'Audit (DO)', 'ProofGen (CSP)', 'Verify (CSP)']
    colors = sns.color_palette("husl", len(y_columns))

    fig, axes = plt.subplots(2, 2, figsize=(18, 14))
    fig.suptitle('Performance Analysis', fontsize=22, fontweight='bold', y=0.97)
    axes = axes.flatten()

    # --- 图 1: 文件大小变化 --- (不变)
    ax1 = axes[0]
    df_filesize.plot(x='Value', y=y_columns, marker='o', ax=ax1, color=colors, ms=8, lw=2.5)
    ax1.set_title('Performance vs. File Size', fontsize=16, fontweight='bold')
    ax1.set_xlabel('File Size (MB)', fontsize=14)
    ax1.set_ylabel('Time (ms)', fontsize=14)
    ax1.legend(legend_labels, fontsize=12)
    ax1.set_yscale('log')
    ax1.tick_params(axis='both', which='major', labelsize=12)

    # --- 图 2: CSP 数量变化 --- (不变)
    ax2 = axes[1]
    df_cspcount.plot(x='Value', y=y_columns, marker='s', ax=ax2, color=colors, ms=8, lw=2.5)
    ax2.set_title('Performance vs. Number of CSPs', fontsize=16, fontweight='bold')
    ax2.set_xlabel('Number of CSPs (K)', fontsize=14)
    ax2.set_ylabel('Time (ms)', fontsize=14)
    ax2.legend(legend_labels, fontsize=12)
    ax2.set_yscale('log')
    ax2.tick_params(axis='both', which='major', labelsize=12)

    # --- 图 3: 数据块数量变化 --- (核心修改)
    ax3 = axes[2]
    df_numblocks.plot(x='Value', y=y_columns, marker='^', ax=ax3, color=colors, ms=8, lw=2.5)
    # ******************** 核心修改点 ********************
    ax3.set_title('Performance vs. Number of Blocks', fontsize=16, fontweight='bold')
    ax3.set_xlabel('Number of Blocks', fontsize=14)
    # ******************************************************
    ax3.set_ylabel('Time (ms)', fontsize=14)
    ax3.legend(legend_labels, fontsize=12)
    ax3.tick_params(axis='both', which='major', labelsize=12)
    
    axes[3].set_visible(False)

    plt.tight_layout(rect=[0, 0, 1, 0.95])
    plt.savefig('performance_plots_final.png', dpi=300, bbox_inches='tight')
    print("Final performance plots saved to 'performance_plots_final.png'")
    plt.show()


if __name__ == '__main__':
    try:
        results_df = pd.read_csv('performance_results.csv')
        results_df.rename(columns={'VerfyTimeMs': 'VerifyTimeMs'}, inplace=True, errors='ignore')
        plot_results(results_df)
    except FileNotFoundError:
        print("Error: 'performance_results.csv' not found. Please run the Go program first.")
    except Exception as e:
        print(f"An error occurred: {e}")